﻿using System;
using Algoritmia.DataStructures.Lists;
using Algoritmia.DataStructures.PriorityQueues;
using Algoritmia.DataStructures.Sets;
using Algoritmia.StateSpaces;
using Algoritmia.Utils.Numerics;

namespace Algoritmia.Schemes.BranchAndBound
{
    public class BestFirstOptimizationSolver<TState,TDecision,TWeight,TSolution>
        where TState : class
        where TWeight : IComparable<TWeight>
    {
        private readonly Func<IForwardStateSpace<TState, TDecision>, IPriorityQueue<BranchAndBoundState<TState,TDecision,TWeight>>> _createPriorityQueue;
        private readonly Func<IForwardStateSpace<TState, TDecision>, ISet<TState>> _createSet;
        private readonly Func<TState, BranchAndBoundDecision<TDecision>, TState, TSolution> _createSolution;
        private readonly TWeight _worst;
        private readonly Func<TWeight, TWeight, TWeight> _opt;
        private static IArithmeticsProvider<TWeight> _arithmeticsProvider = new BuiltinNumericsArithmeticsProvider<TWeight>();

        public BestFirstOptimizationSolver(Func<TState, BranchAndBoundDecision<TDecision>, TState, TSolution> createSolution, 
                                 Func<IForwardStateSpace<TState, TDecision>, IPriorityQueue<BranchAndBoundState<TState, TDecision, TWeight>>> createPriorityQueue = null,
                                 Func<IForwardStateSpace<TState, TDecision>, ISet<TState>> createSet = null)
            : this(_arithmeticsProvider.MaxValue,createSolution,_arithmeticsProvider.Min,createPriorityQueue,createSet)
        {
        }

        public BestFirstOptimizationSolver(TWeight worst, Func<TState, BranchAndBoundDecision<TDecision>, TState, TSolution> createSolution, 
                                 Func<TWeight,TWeight,TWeight> opt,
                                 Func<IForwardStateSpace<TState, TDecision>, IPriorityQueue<BranchAndBoundState<TState, TDecision, TWeight>>> createPriorityQueue = null,
                                 Func<IForwardStateSpace<TState, TDecision>, ISet<TState>> createSet = null)
        {
            _worst = worst;
            _opt = opt;
            _createSolution = createSolution;
            _createPriorityQueue = createPriorityQueue ?? (space => new MaxHeap<BranchAndBoundState<TState, TDecision, TWeight>>());
            _createSet = createSet ?? (space => new DummySet<TState>());
        }

        public static IArithmeticsProvider<TWeight> ArithmeticsProvider
        {
            get { return _arithmeticsProvider; }
            set { _arithmeticsProvider = value; }
        }

        public TSolution Solve(IForwardStateSpace<TState,TDecision> space, Func<TState,TWeight> f,
            Func<TState,TWeight> priorityFunc = null)
        {
            priorityFunc = priorityFunc ?? (s => default(TWeight));
            var open = _createPriorityQueue(space);
            var visited = _createSet(space);
            TWeight bestWeight = _worst;
            TSolution bestSolution = default(TSolution);
            bool hasSolution = false;
            foreach(var s in space.InitialStates)
            {
                open.Add(BranchAndBoundState.Create(s, priorityFunc(s), new BranchAndBoundDecision<TDecision>(), s));
                visited.Add(s);
            }

            while (open.Count > 0)
            {
                var s = open.ExtractOptimal();
                var state = s.Current;
                var prev = s.Previous;
                var decision = s.Decision;
                if ( space.IsFinal(state))
                {
                    var fs = f(state);
                    if ( _opt(fs,bestWeight).CompareTo(bestWeight) != 0)
                    {
                        bestWeight = fs;
                        bestSolution = _createSolution(state,decision,prev);
                        hasSolution = true;
                    }
                }
                foreach(var d in space.Decisions(state))
                {
                    var s1 = space.Decide(state, d);
                    if (!visited.Contains(s1))
                    {
                        open.Add(BranchAndBoundState.Create(s1, priorityFunc(s1), new BranchAndBoundDecision<TDecision>(d, decision), state));
                        visited.Add(s1);
                    }
                }
            }
            if ( !hasSolution)
            {
                throw new InvalidOperationException("solution not found");
            }
            return bestSolution;
        }
    }

    public class BestFirstOptimizationSolver<TState, TDecision, TWeight> : BestFirstOptimizationSolver<TState, TDecision, TWeight, IList<TDecision>>
        where TState : class
        where TWeight : IComparable<TWeight>
    {
        public BestFirstOptimizationSolver(Func<IForwardStateSpace<TState, TDecision>, IPriorityQueue<BranchAndBoundState<TState, TDecision, TWeight>>> createPriorityQueue = null,
                                           Func<IForwardStateSpace<TState, TDecision>, ISet<TState>> createSet = null)
            : base((s, d, p) => d.RecoverDecisions(), createPriorityQueue, createSet)
        {
        }

        public BestFirstOptimizationSolver(TWeight worst, Func<TWeight, TWeight, TWeight> opt,
                              Func<IForwardStateSpace<TState, TDecision>, IPriorityQueue<BranchAndBoundState<TState, TDecision, TWeight>>> createPriorityQueue = null,
                              Func<IForwardStateSpace<TState, TDecision>, ISet<TState>> createSet = null)
            : base(worst, (s, d, p) => d.RecoverDecisions(), opt, createPriorityQueue, createSet)
        {
        }
    }
}
